import argparse

from numpy.core.numeric import fromfunction

from har.dataset import *
from har.data_analysis import *
from har.classifier import *
from util.logger import Recorder

rec_header = "数据集, 分类器, 正确率, 耗时"
rec_file = "output/har/clf_score.csv"
recorder = Recorder(header=rec_header, recorder_file=rec_file, mode="a")


hardataset = HARDATASET("data")
outputdir = "output/har"
sample_rate = 0.01

clfs = ["GBDT", "MLP", "AdaBoost", "XGBoost"]
# clfs = ["AdaBoost"]
classifier = Classifier(*clfs)

"""
PART1:
    data analysis
"""
# train data
def train_data():
    x = hardataset.get_data("train", "x")
    y = hardataset.get_data("train", "y")
    x_, y_ = t_sne(x, y, sample_rate)
    label_dic = hardataset.y2label()

    scatter_with_label(x_, y_, label_dic, title="train-x-y", savefile=f"{outputdir}/train-x-y-tsne.svg")
    sample_mean_line(x, y, label_dic, title="train-x-y-mean", savefile=f"{outputdir}/train-x-y-mean.svg")

# inertial data
def inertial_data():
    inertial_dic, y = get_inertial_data()
    label_dic = hardataset.y2label()

    for key in inertial_dic:
        x_, y_ = t_sne(inertial_dic[key], y, sample_rate)
        scatter_with_label(x_, y_, label_dic, title=key, savefile=f"{outputdir}/{key}-tsne.svg")
        sample_mean_line(inertial_dic[key], y, label_dic, title=key, savefile=f"{outputdir}/{key}-mean.svg")

def get_inertial_data(train=True):
    meta = "train" if train else "test"
    header = [meta, "Inertial Signals"]
    parts = ["body_acc_x", "body_acc_y", "body_acc_z", "body_gyro_x", "body_gyro_y", "body_gyro_z", "total_acc_x", "total_acc_y", "total_acc_z"]
    inertial_dic = {}
    for p in parts:
        data = hardataset.get_data(*header, p)
        inertial_dic[p] = data

    y = hardataset.get_data(meta, "y")
    return inertial_dic, y

"""
PART2:
    classifier train
"""
# clf input: train file
def train_clf():
    x = hardataset.get_data("train", "x")
    y = hardataset.get_data("train", "y")
    test_x = hardataset.get_data("test", "x")
    test_y = hardataset.get_data("test", "y")

    elapsed = classifier.fit(x, y)
    acc_dict = classifier.score(test_x, test_y)
    for clf in clfs:
            recorder.record_one(f"train-x, {clf}, {acc_dict[clf]}, {elapsed[clf]}")   

# clf input: inertial file
def inertial_clf():
    inertial_dic, y = get_inertial_data(train=True)
    innertial_test_dic, test_y = get_inertial_data(train=False)

    for key in inertial_dic:
        print(key)
        elapsed = classifier.fit(inertial_dic[key], y)
        acc_dict = classifier.score(innertial_test_dic[key], test_y) 
        for clf in clfs:
            recorder.record_one(f"inertial-{key}, {clf}, {acc_dict[clf]}, {elapsed[clf]}")   

arg_cli = argparse.ArgumentParser()
arg_cli.add_argument("func", help="train data analysis", type=str)

args = arg_cli.parse_args()
if __name__ == "__main__":
    func_map = {
        "td": train_data,
        "id": inertial_data,
        "tc": train_clf,
        "ic": inertial_clf,
    }

    func = args.func
    if func not in func_map:
        print("func:")
        [print(f"{key}: {func_map[key].__name__}") for key in func_map]
        exit()
    func_map[func]()
